import asyncio
import json
from typing import List, Dict, Optional
from jinja2 import Environment, FileSystemLoader
import json_repair
from research_agent.core.query import Query
from research_agent.core.config import Config
from research_agent.core.general_llm import LLM
import re
from typing import Union, Optional
from pyaml_env import parse_config
from pathlib import Path

class AddCitation:
    """
    用于为研究陈述添加引用的核心类。
    
    该类负责：
    1. 从研究陈述中提取关键信息
    2. 查询相关文献片段
    3. 使用LLM生成适当的引用
    4. 处理批量陈述的引用添加
    """
    def __init__(self):
        """
        初始化AddCitation实例。
        
        加载配置文件，初始化查询和LLM实例，
        并加载引用提示模板。
        """
        configs = parse_config(Config.YAML_CONFIG)
        self.query = Query()
        self.llm = LLM(config=configs[Config.DEFAULT_MODEL])
        self.top_k = 5
        
        # 获取当前文件所在目录的路径
        base_path = Path(__file__).parent / "prompts"

        # 修改文件路径的获取方式
        add_citations_prompt_file = base_path / "add_citations.jinja"
        with open(add_citations_prompt_file, "r",encoding="utf-8") as f:
            self.add_citations_prompt_template = Environment().from_string(f.read())


    async def _retrieve_context(self, statement: Union[str, dict]) -> Optional[str]:
        """
        检索与陈述相关的上下文片段。
        
        参数:
            statement: 需要引用的陈述，可以是字符串或包含未支持论文信息的字典
            
        返回:
            包含相关文献片段的字符串，如果没有找到相关片段则返回None
        """
        context = ""
        if isinstance(statement, dict) and "unsupported_papers" in statement:
            unsupported_papers = [json_repair.loads(re.findall(
                r"<sup>(.*?)</sup>", s)[0]) for s in statement["unsupported_papers"]]
            chunks = await self.query.query_by_content(statement["statement_explanation"], top_k=self.top_k+len(statement["unsupported_papers"]))
            chunk_text = []
            for chunk in chunks:
                if (chunk["entity"]["paper_id"] in [p["paper_id"] for p in unsupported_papers]) and (str(chunk["entity"]["chunk_id"]) in [str(p["chunk_id"]) for p in unsupported_papers]):
                    continue
                else:
                    formatted_text = f"paper_title:{chunk['entity']['paper_title']}\n" \
                        f"paper_id:{chunk['entity']['paper_id']}\n" \
                        f"chunk_id:{chunk['entity']['chunk_id']}\n" \
                        f"chunk_text:{chunk['entity']['chunk_text']}"
                    chunk_text.append(formatted_text)
            context = "\n".join(chunk_text)
        else:
            chunks = await self.query.query_by_content(statement, top_k=self.top_k)
            context = "\n".join(
                f"""paper_title:{chunk['entity']['paper_title']}
                                paper_id:{chunk['entity']['paper_id']}
                                chunk_id:{chunk['entity']['chunk_id']}
                                chunk_text:{chunk['entity']['chunk_text']}"""
                for chunk in chunks
            )
        if not chunks:
            return None
        return chunks

    def _prepare_messages(self, statement: str, context: str) -> List[Dict]:
        """
        准备LLM的输入消息。
        
        参数:
            statement: 需要引用的陈述
            context: 检索到的相关文献片段
            
        返回:
            包含系统提示和用户提示的消息列表
        """
        system_prompt = self.add_citations_prompt_template.render(
            role="system")
        user_prompt = self.add_citations_prompt_template.render(
            role="user", statement=statement, retrieved_chunk=context
        )
        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]

    async def _process_statement(self, statement: Union[str, dict]) -> Dict:
        """
        处理单个陈述并生成引用。
        
        参数:
            statement: 需要处理的陈述，可以是字符串或字典
            
        返回:
            包含处理结果的字典，格式为：
            {
                "statement": 原始陈述,
                "answer": LLM生成的回答,
                "related_papers": 相关论文列表
            }
            如果处理失败，则返回包含错误信息的字典
        """
        try:
            if isinstance(statement, dict):
                unsupported_papers = [json_repair.loads(re.findall(
                    r"<sup>(.*?)</sup>", s)[0]) for s in statement["unsupported_papers"]]
                chunks = await self.query.query_by_content(statement["statement"], top_k=self.top_k+len(statement["unsupported_papers"]))
                chunk_text = []
                for chunk in chunks:
                    if (chunk["entity"]["paper_id"] in [p["paper_id"] for p in unsupported_papers]) and (str(chunk["entity"]["chunk_id"]) in [str(p["chunk_id"]) for p in unsupported_papers]):
                        continue
                    else:
                        formatted_text = f"paper_title:{chunk['entity']['paper_title']}\n" \
                            f"paper_id:{chunk['entity']['paper_id']}\n" \
                            f"chunk_id:{chunk['entity']['chunk_id']}\n" \
                            f"chunk_text:{chunk['entity']['chunk_text']}"
                        chunk_text.append(formatted_text)
                context = "\n".join(chunk_text)
            if isinstance(statement, str):
                chunks = await self.query.query_by_content(statement, self.top_k)
                context = "\n".join(
                    f"""paper_title:{chunk['entity']['paper_title']}
                                paper_id:{chunk['entity']['paper_id']}
                                chunk_id:{chunk['entity']['chunk_id']}
                                chunk_text:{chunk['entity']['chunk_text']}"""
                    for chunk in chunks
                )
            if not context:
                return {"statement": statement, "answer": "No relevant information found.", "related_papers": []}

            messages = self._prepare_messages(statement, context)
            response = await self.llm.completion(messages)
            return json_repair.loads(response)
        except Exception as e:
            return {
                "statement": statement,
                "error": str(e),
                "related_papers": []
            }

    async def add_citations(self, statements: List[str]) -> List[Dict]:
        """
        批量添加引用的主入口点。
        
        参数:
            statements: 需要添加引用的陈述列表
            
        返回:
            包含每个陈述处理结果的字典列表
            
        异常:
            ValueError: 如果输入不是非空列表
        """
        if not isinstance(statements, list) or not statements:
            raise ValueError("Input must be a non-empty list of statements")

        tasks = [self._process_statement(stmt) for stmt in statements]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        return [
            result for result in results
            if not isinstance(result, Exception)
        ]


async def main():
    """
    主函数，用于测试AddCitation类的功能。
    
    示例:
        1. 创建AddCitation实例
        2. 定义测试陈述
        3. 调用add_citations方法
        4. 打印结果
    """
    add_citation = AddCitation()

    statements = [
        "Deep learning has shown remarkable success in computer vision tasks.",
        "The transformer architecture revolutionized natural language processing."
    ]

    results = await add_citation.add_citations(statements)
    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    asyncio.run(main())
